import torch
from sentence_transformers import SentenceTransformer

# Load a model to finetune with
embedding_model = SentenceTransformer(model_name_or_path="../transformers/fin-bge-m3",
                                      local_files_only=True)

# 输入query
query = ['优化MySQL连接池等待时间']

documents = [
    "优化方案：\n1. 使用G1垃圾回收器替代CMS：JVM参数添加 `-XX:+UseG1GC`\n2. 调整最大堆大小：`-Xmx4g -Xms4g` 避免动态扩容\n3. 优化对象分配：减少大对象直接进入老年代\n验证效果：Full GC频率从每小时10次降至0.5次",
    "无关方案：\n增加线程池大小到200个线程\n效果：CPU使用率上升但GC问题未改善",
    "优化步骤：\n1. 调整Druid连接池配置：\n   - `maxActive=50` (原值20)\n   - `maxWait=1000ms` (原值3000ms)\n2. 添加连接有效性检查：`testWhileIdle=true`\n效果：平均等待时间从2.3s降至200ms",
    "无效尝试：\n将`synchronized`改为`ReentrantLock`\n实测：切换开销无显著变化"
]

query_result = embedding_model.encode(query, normalize_embeddings=True)

document_result = embedding_model.encode(documents, normalize_embeddings=True)

# 计算相似度
similarity_scores = embedding_model.similarity(query_result, document_result)[0]
scores, indices = torch.topk(similarity_scores, k=2)

print("\nQuery:", query)
print("Top 5 most similar sentences in corpus:")

for score, idx in zip(scores, indices):
    print(f"(Score: {score:.4f})", documents[idx])
